[https://nvbugs/6902355][fix] Keep TopK rows in the DSA sparse MLA test KV pool by lfr-0531 · Pull Request #20283 · NVIDIA/TensorRT-LLM

lfr-0531 · 2026-10-14T09:45:17Z

Description

Fixes the intermittent illegal memory access in
unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py::test_sparse_attention_mla reported in
https://nvbugs/6902355 (LLM/main/L0_PostMerge #3009, DGX_B300-4_GPUs-PyTorch-Post-Merge-2, FP8 KV,
context_sequence_lengths=[10]). It recurs in CI pipelines (see CI impact), but the repair bot
could not reproduce it standalone (75/75 passes).

Root cause. The test sizes its KV cache to the blocks the short sequences need. With
context_sequence_lengths=[10] that is a single 64-token block, so the primary pool is 36,864 B with
FP8 KV and 73,728 B with BF16 KV. The trtllm-gen static-token-sparse context FMHA used by DSA on
SM100f/SM103 can load KV rows beyond the selected ones. Its token-sparse K TMA descriptor has
INT_MAX rows (KernelParams.h), and the dispatcher sets mNumPagesInMemPool = INT_MAX, so nothing
clips those loads to the pool. Whether a load past the end faults depends only on whether the
address range after the pool happens to be mapped. That changes with the allocation history of a
long pytest process. That is why it fails in CI, for example about 1,430 tests into the #3009
process, but not standalone.

Fix. Keep at least SPARSE_TOPK (2048) rows in the test's KV pool. Every real deployment has
far larger pools (blocks × layers × 64 rows), and 2048 rows is the same TopK-row invariant the DSA
backend already enforces for its NVFP4 context scratch.

CI impact. The fault is not a one-off. In the last 30 days (2026-09-14 to 10-14), 40 CI runs
failed first on an FP8 context_sequence_lengths: [10] case with this exact signature:
test_dsa_sparse_mla.py:154 IMA, then cudaEventCreate IMA in BlockManager::releasePools.

By stage: 30 DGX_B200-PyTorch, 9 B300-PyTorch and 1 DGX_B300 post-merge run (feat: Update logits bitmask kernel to v3 #3009).
32 of the 40 hit the generation_seq_len_q: 4 case and 8 the generation_seq_len_q: 1 case.
Each poisons the CUDA context, so every later test in the same pytest process fails on setup
with an illegal memory access. For example, all 18 setup failures of
test_sparse_mla_forward.py::test_forward_sparse_mla_unified[deepseek_v4-auto-small_prefill] in
that window come from these runs.
Most of these runs were triaged as PR-induced or intermittent.

Evidence (B300 / GB110, CI wheels of post-merge #3009 ec357dd9 and #3018 af469207):

Deterministic repro. A diagnostic pytest plugin copied each KV pool into a CUDA VMM buffer
that ends right before reserved but unmapped VA, and rewrote kv_cache_pool_pointers.
- The ctx=[10] cases then fail on every run, on both builds, with CUDA_ERROR_ILLEGAL_ADDRESS at
  the FMHA cuLaunchKernelEx (CUDA_LAUNCH_BLOCKING=1).
- The same layout passes with ≥ 2 MB of mapped slack after the pool, or with the unmapped range
  placed before the pool.
Faulting kernel. compute-sanitizer memcheck on the guarded pool reports every error as
Warp illegal address in
fmhaSm100fKernel_QkvE4m3OBfloat16HQk576HV512HVPerCta256PagedKvDenseStaticTokenSparseP1VarSeqQ64Kv128Static2CtaKeepsAbForGen.
The BF16 run reports the QkvBfloat16 variant of the same kernel.
Scope.
- For a 10-token request at the start of a 1-block pool, the kernel touches about 57–62 rows past
  the end (FP8).
- Faults appear only with pools of ≤ 2 blocks (128 rows).
- Pools of 3, 4, 8, 65, 66 and 80 blocks pass with the request in the first, middle or last block,
  for ctx 10/100/1000 and FP8/BF16.
Fix validation.
- Under the same guard, the patched test passes all 12 cases on both builds.
- It also passes without the guard.

A kernel-side follow-up for the trtllm-gen owners is to bound the token-sparse KV descriptor rows to
the real pool, so short or tail requests can never load past the allocation. This PR only removes the
test's unrealistic 1-block pool.

Test Coverage

unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py, run in the
unittest/_torch/attention entries of l0_b200, l0_b300 and l0_dgx_b300.

PR Checklist

Please check this after reviewing the above items as appropriate for this PR.

🤖 Generated with Claude Code

yuxianq · 2026-10-14T09:52:35Z

@yunruis, could you help to review this PR as trt-llm-torch-attention-devs? Thanks~

Assignment explanation

Components	integration/attention-tests (100%)
Tags	algorithm:sparse (100%); specific-algorithm:dsa (100%); representation:mla (100%); phase:prefill (100%); features:paged-kv (100%)
PR effort	3 points — The capacity floor is localized, but validating it requires checking block sizing and the claimed relationship between TopK and kernel overfetch. That memory-access assumption raises review effort beyond a mechanical test edit.
Affinity (0–1)	0.468 component/tag match
Outstanding points at selection	56 → 56 (+0)
Points already reserved on this PR	3
Effective weight	1
Completed review points (preceding 168 hours)	16
Combined points after assignment	72
Demand-derived target	72.00 points (soft share)
Additional fair share before assignment	0.00 points
Share explanation	Within demand-derived share.
Selection rule	balance-first maintenance

Affinity and workload are snapshots at selection, not live counters. Existing reservations on this PR are counted once.

coderabbitai · 2026-10-14T09:55:29Z

Important

Review skipped

Review was skipped as selected files did not have any reviewable changes.

⚙️ Run configuration

Configuration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5cbfad93-4231-4c8f-909e-8aa0f33c0e07

📥 Commits

Reviewing files that changed from the base of the PR and between c56d2ce and 45acc45.

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f0dfc4fd-4ab6-4f35-a3af-14cf4bb83791

📥 Commits

Reviewing files that changed from the base of the PR and between 30dc23f and 582accb.

📒 Files selected for processing: 1

tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

Walkthrough

The sparse MLA test now ensures that the KV cache capacity is at least SPARSE_TOPK tokens after calculating capacity from sequence lengths.

Changes

Sparse MLA test

Layer / File(s)	Summary
KV cache capacity adjustment `tests/unittest/_torch/attention/sparse/dsa/test_dsa_sparse_mla.py`	The test raises the calculated KV cache capacity to at least `SPARSE_TOPK` tokens.

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~8 minutes

Change: Bug fix

Suggested reviewers: bowenfu

Merge Risk: ⚪ Minimal · up to 582ac

This test-only change enlarges the sparse MLA backing pool without changing logical sequence lengths. No actionable merge-blocking risk is established; merge risk is minimal, subject to normal checks.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 1 functions across 1 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Title check		The title clearly identifies the NVBugs issue, uses the valid fix type, and summarizes the primary change to the DSA sparse MLA test KV pool.
Description check		The description includes the required Description, Test Coverage, and PR Checklist sections. It clearly explains the failure, root cause, fix, validation, and affected tests.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create a new PR

Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

_{Comment @coderabbitai help to get the list of available commands.}

lfr-0531 · 2026-10-14T13:28:33Z

/bot run --extra-stage "DGX_B300-4_GPUs-PyTorch-Post-Merge-1,DGX_B300-4_GPUs-PyTorch-Post-Merge-2"

tensorrt-cicd · 2026-10-14T13:37:32Z

PR_Github #77816 [ run ] triggered by Bot. Commit: 582accb Link to invocation

tensorrt-cicd · 2026-10-14T15:49:36Z

PR_Github #77816 [ run ] completed with state SUCCESS. Commit: 582accb
/LLM/main/L0_MergeRequest_PR pipeline #64149 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

lfr-0531 · 2026-10-14T15:56:21Z

/bot run --extra-stage "DGX_B300-4_GPUs-PyTorch-Post-Merge-1,DGX_B300-4_GPUs-PyTorch-Post-Merge-2"

tensorrt-cicd · 2026-10-14T16:04:42Z

PR_Github #77860 [ run ] triggered by Bot. Commit: c56d2ce Link to invocation

tensorrt-cicd · 2026-10-14T18:37:17Z

PR_Github #77860 [ run ] completed with state SUCCESS. Commit: c56d2ce
/LLM/main/L0_MergeRequest_PR pipeline #64185 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

lfr-0531 · 2026-10-14T18:42:15Z

/bot run --extra-stage "DGX_B300-4_GPUs-PyTorch-Post-Merge-1,DGX_B300-4_GPUs-PyTorch-Post-Merge-2"

tensorrt-cicd · 2026-10-14T18:51:28Z

PR_Github #77897 [ run ] triggered by Bot. Commit: c56d2ce Link to invocation

tensorrt-cicd · 2026-10-14T23:24:01Z

PR_Github #77897 [ run ] completed with state SUCCESS. Commit: c56d2ce
/LLM/main/L0_MergeRequest_PR pipeline #64210 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

lfr-0531 · 2026-10-14T23:29:14Z

/bot run --extra-stage "DGX_B300-4_GPUs-PyTorch-Post-Merge-1,DGX_B300-4_GPUs-PyTorch-Post-Merge-2"

tensorrt-cicd · 2026-10-14T23:36:07Z

PR_Github #77926 [ run ] triggered by Bot. Commit: c56d2ce Link to invocation

tensorrt-cicd · 2026-10-15T02:16:19Z

PR_Github #77926 [ run ] completed with state FAILURE. Commit: c56d2ce
/LLM/main/L0_MergeRequest_PR pipeline #64234 completed with status: 'UNSTABLE'

CI Report

⚠️ Multi-GPU Label Required:
Multi-GPU tests require the ci: full pre-merge approved label on this PR. Either:

Wait for the PR to be fully approved — the label is added automatically once approval is complete. Having unresolved open comments is fine, or
If needed, ask a member of NVIDIA/trt-llm-ci-approvers to add the label manually.
Then re-trigger CI with the same bot command (no rebase needed).

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

lfr-0531 · 2026-10-15T08:31:07Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-10-15T08:38:02Z

PR_Github #78069 [ run ] triggered by Bot. Commit: c56d2ce Link to invocation

tensorrt-cicd · 2026-10-15T10:00:53Z

PR_Github #78069 [ run ] completed with state SUCCESS. Commit: c56d2ce
/LLM/main/L0_MergeRequest_PR pipeline #64348 completed with status: 'SUCCESS'

CI Report

Link to invocation

github-actions · 2026-10-16T03:31:17Z

Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements.

lfr-0531 · 2026-10-16T03:31:32Z

/bot run --extra-stage "DGX_B300-4_GPUs-PyTorch-Post-Merge-1,DGX_B300-4_GPUs-PyTorch-Post-Merge-2"

tensorrt-cicd · 2026-10-16T03:38:08Z

PR_Github #78336 [ run ] triggered by Bot. Commit: c56d2ce Link to invocation

yuxianq · 2026-10-16T04:00:54Z

@yuxianq, thanks for picking up this PR! I’ve handed off the trt-llm-torch-attention-devs review from @yunruis to you.

The 3 attention-review points now belong to @yuxianq; @yunruis’s attention reservation has been released. Your review is already submitted, so the points count as completed work; no new review request was added.

@yunruis, no further attention review is requested from you. Thanks~

tensorrt-cicd · 2026-10-16T06:45:22Z

PR_Github #78336 [ run ] completed with state FAILURE. Commit: c56d2ce
/LLM/main/L0_MergeRequest_PR pipeline #64569 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

…st KV pool test_sparse_attention_mla sized its KV cache to the blocks the short sequences need, so context_sequence_lengths=[10] got a single 64-token block. The static token-sparse context FMHA can load KV rows beyond the selected ones, and its TMA descriptor does not bound rows to the pool, so such a pool can be read past its end and faults whenever the following address range is unmapped. Keep at least TopK rows in the pool, as any real deployment does. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com>

lfr-0531 · 2026-10-16T06:48:58Z

/bot run --extra-stage "DGX_B300-4_GPUs-PyTorch-Post-Merge-1,DGX_B300-4_GPUs-PyTorch-Post-Merge-2"

tensorrt-cicd · 2026-10-16T06:55:12Z

PR_Github #78398 [ run ] triggered by Bot. Commit: 45acc45 Link to invocation

tensorrt-cicd · 2026-10-16T07:53:47Z

PR_Github #78398 [ run ] completed with state FAILURE. Commit: 45acc45
/LLM/main/L0_MergeRequest_PR pipeline #64622 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

lfr-0531 · 2026-10-16T07:54:24Z

/bot run --extra-stage "DGX_B300-4_GPUs-PyTorch-Post-Merge-1,DGX_B300-4_GPUs-PyTorch-Post-Merge-2"

tensorrt-cicd · 2026-10-16T08:01:55Z

PR_Github #78417 [ run ] triggered by Bot. Commit: 45acc45 Link to invocation

tensorrt-cicd · 2026-10-16T08:24:09Z

PR_Github #78417 [ run ] completed with state FAILURE. Commit: 45acc45
/LLM/main/L0_MergeRequest_PR pipeline #64638 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

lfr-0531 · 2026-10-16T08:24:45Z

/bot run --extra-stage "DGX_B300-4_GPUs-PyTorch-Post-Merge-1,DGX_B300-4_GPUs-PyTorch-Post-Merge-2"

tensorrt-cicd · 2026-10-16T08:31:07Z

PR_Github #78431 [ run ] triggered by Bot. Commit: 45acc45 Link to invocation

tensorrt-cicd · 2026-10-16T11:44:02Z

PR_Github #78431 [ run ] completed with state SUCCESS. Commit: 45acc45
/LLM/main/L0_MergeRequest_PR pipeline #64649 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

lfr-0531 · 2026-10-16T15:18:36Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-10-16T15:25:49Z

PR_Github #78504 [ run ] triggered by Bot. Commit: 45acc45 Link to invocation

tensorrt-cicd · 2026-10-16T16:12:36Z

PR_Github #78504 [ run ] completed with state SUCCESS. Commit: 45acc45
/LLM/main/L0_MergeRequest_PR pipeline #64704 completed with status: 'SUCCESS'

CI Report

Link to invocation

lfr-0531 requested a review from a team as a code owner October 14, 2026 09:45

lfr-0531 requested review from farazkh80 and yunruis October 14, 2026 09:45

github-actions Bot assigned lfr-0531 Oct 14, 2026

yuxianq removed the request for review from farazkh80 October 14, 2026 09:52

lfr-0531 force-pushed the user/fanrongl/fix-nvbug-6902355-dsa-test-kv-pool branch from 582accb to c56d2ce Compare October 14, 2026 15:55

lfr-0531 requested review from yuxianq and removed request for yunruis October 16, 2026 02:22

trtllm-agent mentioned this pull request Oct 16, 2026

[https://nvbugs/6930163][fix] Shut down Inductor compile workers in KV pool rebalance accuracy test #20450

Closed

2 tasks done

yuxianq approved these changes Oct 16, 2026

View reviewed changes

trtllm-agent added the ci: full pre-merge approved label Oct 16, 2026

lfr-0531 enabled auto-merge (squash) October 16, 2026 05:52

lfr-0531 force-pushed the user/fanrongl/fix-nvbug-6902355-dsa-test-kv-pool branch from c56d2ce to 45acc45 Compare October 16, 2026 06:48

lfr-0531 merged commit b519164 into NVIDIA:main Oct 16, 2026

Conversation

lfr-0531 commented Oct 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

PR Checklist

Uh oh!

yuxianq commented Oct 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

coderabbitai Bot commented Oct 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review skipped

Walkthrough

Changes

Uh oh!

lfr-0531 commented Oct 14, 2026

Uh oh!

tensorrt-cicd commented Oct 14, 2026

Uh oh!

tensorrt-cicd commented Oct 14, 2026

Uh oh!

lfr-0531 commented Oct 14, 2026

Uh oh!

tensorrt-cicd commented Oct 14, 2026

Uh oh!

tensorrt-cicd commented Oct 14, 2026

Uh oh!

lfr-0531 commented Oct 14, 2026

Uh oh!

tensorrt-cicd commented Oct 14, 2026

Uh oh!

tensorrt-cicd commented Oct 14, 2026

Uh oh!

lfr-0531 commented Oct 14, 2026

Uh oh!

tensorrt-cicd commented Oct 14, 2026

Uh oh!

tensorrt-cicd commented Oct 15, 2026

Uh oh!

lfr-0531 commented Oct 15, 2026

Uh oh!

tensorrt-cicd commented Oct 15, 2026

Uh oh!

tensorrt-cicd commented Oct 15, 2026

Uh oh!

github-actions Bot commented Oct 16, 2026

Uh oh!

lfr-0531 commented Oct 16, 2026

Uh oh!

tensorrt-cicd commented Oct 16, 2026

Uh oh!

yuxianq commented Oct 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tensorrt-cicd commented Oct 16, 2026

Uh oh!

lfr-0531 commented Oct 16, 2026

Uh oh!

tensorrt-cicd commented Oct 16, 2026

Uh oh!

tensorrt-cicd commented Oct 16, 2026

Uh oh!

lfr-0531 commented Oct 16, 2026

Uh oh!

tensorrt-cicd commented Oct 16, 2026

Uh oh!

tensorrt-cicd commented Oct 16, 2026

Uh oh!

lfr-0531 commented Oct 16, 2026

Uh oh!

tensorrt-cicd commented Oct 16, 2026

Uh oh!

tensorrt-cicd commented Oct 16, 2026

Uh oh!

lfr-0531 commented Oct 16, 2026

Uh oh!

tensorrt-cicd commented Oct 16, 2026

Uh oh!

tensorrt-cicd commented Oct 16, 2026

lfr-0531 commented Oct 14, 2026 •

edited

Loading

yuxianq commented Oct 14, 2026 •

edited

Loading

coderabbitai Bot commented Oct 14, 2026 •

edited

Loading

yuxianq commented Oct 16, 2026 •

edited

Loading